# -*- coding: utf-8 -*-
# @Time  : 2021/3/21 2:08
# @Author : zhoujiangtao
# @Desc : ==============================================
# Life is Short I Use Python!!!                      
# If this runs wrong,don't ask me,I don't know why;  
# If this runs right,thank god,and I don't know why. 
# Maybe the answer,my friend,is blowing in the wind. 
# ======================================================
import os

import cv2

import data
from data import labels
from train import Net
import torch
import torchvision


def get_label_char(index: int) -> chr:
    return chr(labels[index])


def showResult(max_index, label):
    real_lst: list = []
    check_lst: list = []
    for i in range(0, len(max_index)):
        real_index = label[i][0].item()
        real_label = get_label_char(real_index)
        real_lst.append(real_label)
        # real_lst.append(f"{real_index}:{real_label}")

        check_index = max_index[i][0].item()
        check_label = get_label_char(check_index)
        check_lst.append(check_label)
        # check_lst.append(f"{check_index}:{check_label}")

    print("真实标签：", "".join(real_lst))
    print("识别结果：", "".join(check_lst), "\n")
    # print("真实标签：", "\t".join(real_lst))
    # print("识别结果：", "\t".join(check_lst))


def _test():
    net = Net()
    sd = torch.load("./model/net.pt")
    net.load_state_dict(sd)
    dl = data.testloader()
    correct = 0
    total = 0
    for step, d in enumerate(dl):
        out = net(d["data"])
        max_index = torch.argmax(out, keepdim=True, dim=1)
        label = d["label"].reshape(-1, 1)
        eq = label.eq(max_index).int()
        correct += eq.sum().item()
        total += dl.batch_size

        showResult(max_index, label)

        img = cv2.imread("data/image_test/2/2d2w._2.png")
        img = torchvision.utils.make_grid(d["data"])
        # 旋转90°
        img = img.numpy().transpose(1, 2, 0)
        cv2.imshow('win', img)
        # cv2.imshow('win', img)  # opencv显示需要识别的数据图片
        cv2.waitKey(0)

    print("{}% on {} images test".format(correct * 100 / total, total))


if __name__ == "__main__":
    _test()
